import joblib
import gradio as gr
import numpy as np
from PIL import Image

# 加载保存的KNN模型
best_knn_model = joblib.load('best_knn_model.pickle')

def predict_digit(drawing):
    try:
        # 获取PIL 图像
        image = drawing['composite']

        # 调整图像大小为 10x10
        image = image.resize((8, 8), Image.Resampling.LANCZOS)

        # 将 PIL 图像转换为 NumPy 数组
        image_array = np.array(image)

        # 转换为灰度图像
        if image_array.ndim == 3 and image_array.shape[-1] == 3:
            image_array = np.dot(image_array[..., :3], [0.2989, 0.5870, 0.1140])
        elif image_array.ndim == 3 and image_array.shape[-1] == 4:
            image_array = np.dot(image_array[..., :3], [0.2989, 0.5870, 0.1140])

        # 归一化
        image_array = image_array / 255.0

        # 展平图像
        image_array = image_array.reshape((1, -1))

        # 进行预测并返回结果
        prediction = best_knn_model.predict(image_array)[0]
        return str(prediction)
    except Exception as e:
        return f"Error: {e}"

def create_gradio_interface():
    inp = gr.Sketchpad(label="Draw a digit", type='pil')  # 创建 Sketchpad 组件并指定类型
    iface = gr.Interface(fn=predict_digit,
                         inputs=inp,
                         outputs="text",
                         live=False)

    return iface

if __name__ == "__main__":
    demo = create_gradio_interface()
    demo.launch()

